'''
Author: cw
Date: 2025-08-18
Description: ...
'''
import asyncio

from utils import http_util
from modules.data_models import ProcessTask
from modules.image_base_provider import ImageBaseProvider


class ImageMotaProvider(ImageBaseProvider):
    """本地 生成图片"""
    base_url = 'https://api-inference.modelscope.cn/'
    api_key = "ms-04c601b8-2489-4b24-931e-28454f47b8fc"
    common_headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    async def create(self, task: ProcessTask):
        '''create'''
        data = {
            "model": task.model,
            "prompt": task.prompt,
            'n': task.num,
            'size': task.size
        }
        if task.extra:
            data.update(task.extra)
        res = await http_util.aio_post(
            f"{self.base_url}v1/images/generations",
            json=data,
            headers={**self.common_headers, "X-ModelScope-Async-Mode": "true"},
            timeout=60
        )
        if not (res and res.get('task_id')):
            raise Exception(f"Image Generation Failed.res={res}")
        return res.get('task_id')

    async def info(self, task_id):
        """get info"""
        while True:
            res = await http_util.aio_get(
                f"{self.base_url}v1/tasks/{task_id}",
                headers={**self.common_headers, "X-ModelScope-Task-Type": "image_generation"},
                timeout=60
            )
            print(res)
            if not (res and res.get("task_status")):
                raise Exception(f"Image Generation Failed.res={res}")

            if res["task_status"] == "SUCCEED":
                # image_content = requests.get(res["output_images"][0], timeout=30).content
                # image = Image.open(BytesIO(image_content))
                # image.save("result_image.jpg")
                # return base64.b64encode(image_content).decode("utf-8")
                return res["output_images"]

            if res["task_status"] == "FAILED":
                raise Exception(f"Image Generation Failed.res={res}")

            await asyncio.sleep(3)

    async def run(self, task: ProcessTask):
        """run"""
        task_id = await self.create(task)
        return await self.info(task_id)
